# domain_analysis.py
import cv2
import numpy as np
import matplotlib.pyplot as plt
import os
from glob import glob


def analyze_domain_differences(rural_dir, urban_dir, model):
    """
    分析农村和城市数据的域差异
    """
    print("=== 域差异分析 ===")

    # 分析图像统计特征
    rural_images = glob(os.path.join(rural_dir, '*.png'))[:5]  # 抽样分析
    urban_images = glob(os.path.join(urban_dir, '*.png'))[:5]

    rural_stats = calculate_image_stats(rural_images)
    urban_stats = calculate_image_stats(urban_images)

    print("农村图像统计:")
    for key, value in rural_stats.items():
        print(f"  {key}: {value:.3f}")

    print("城市图像统计:")
    for key, value in urban_stats.items():
        print(f"  {key}: {value:.3f}")

    # 可视化对比
    visualize_domain_comparison(rural_images[0], urban_images[0])


def calculate_image_stats(image_paths):
    """计算图像统计特征"""
    stats = {
        'mean_intensity': 0,
        'std_intensity': 0,
        'contrast': 0,
        'entropy': 0
    }

    for path in image_paths:
        img = cv2.imread(path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        stats['mean_intensity'] += np.mean(img)
        stats['std_intensity'] += np.std(img)
        stats['contrast'] += np.max(img) - np.min(img)
        stats['entropy'] += calculate_entropy(img)

    # 取平均值
    for key in stats:
        stats[key] /= len(image_paths)

    return stats


def calculate_entropy(image):
    """计算图像熵"""
    hist = cv2.calcHist([image], [0], None, [256], [0, 256])
    hist = hist / hist.sum()
    entropy = -np.sum(hist * np.log2(hist + 1e-10))
    return entropy


def visualize_domain_comparison(rural_path, urban_path):
    """可视化域差异"""
    rural_img = cv2.imread(rural_path)
    rural_img = cv2.cvtColor(rural_img, cv2.COLOR_BGR2RGB)

    urban_img = cv2.imread(urban_path)
    urban_img = cv2.cvtColor(urban_img, cv2.COLOR_BGR2RGB)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    ax1.imshow(rural_img)
    ax1.set_title('农村图像示例')
    ax1.axis('off')

    ax2.imshow(urban_img)
    ax2.set_title('城市图像示例')
    ax2.axis('off')

    plt.tight_layout()
    plt.savefig('domain_comparison.png', dpi=300, bbox_inches='tight')
    plt.show()


if __name__ == '__main__':
    rural_dir = 'inputs/Test/Rural/images_png'
    urban_dir = 'inputs/Test/Urban/images_png'

    analyze_domain_differences(rural_dir, urban_dir, None)